'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: token_buffer_memory.py
* @Time: 2025/11/17
* @All Rights Reserve By Brtc
'''
from dataclasses import dataclass

from langchain_core.language_models import BaseLanguageModel
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, trim_messages, get_buffer_string
from sqlalchemy.sql.expression import desc

from internal.entity.conversation_entity import MessageStatus
from internal.model import Conversation, Message
from pkg.sqlalchemy import SQLAlchemy


@dataclass
class TokenBufferMemory:
    """基于token 计数的缓冲记忆组件"""
    db:SQLAlchemy
    conversation:Conversation
    model_instance:BaseLanguageModel

    def get_history_prompt_message(self,
                                   max_token_limit:int =2000,
                                   massage_limit :int=10)->list[AnyMessage]:
        """根据传递的token限制 + 消息条数 限制获取指定会话模型的历史消息列表"""
        #1、判断会话模型是否存在， 如果不存在则直接返回空列表
        if self.conversation is None:
            return []
        #2、查询该会话的消息列表，并且使用时间进行倒序排序,匹配会话id
        messages = self.db.session.query(Message).filter(
            Message.conversation_id==self.conversation.id,
            Message.answer !="",
            Message.is_deleted == False,
            Message.status == MessageStatus.NORMAL
        ).order_by(desc("created_at")).limit(massage_limit).all()
        messages = list(reversed(messages))
        #3、将message 转换成Langchain 的消息列表
        prompt_messages = []

        for message in messages:
            prompt_messages.extend([
                HumanMessage(content=message.query),
                AIMessage(content=message.answer),
            ])

        return trim_messages(
            messages=prompt_messages,
            max_token=max_token_limit,
            token_counter = self.model_instance,
            strategy = "last"
        )


    def get_history_prompt_text(self,
                                human_prefix = "Human",
                                ai_prefix = "Ai",
                                max_token_limit:int=2000,
                                massage_limit :int=10)->str:
        """根据传递的数据获取指定会话的历史消息(短期记忆的轮数 + 长期记忆的摘要)"""
        messages = self.get_history_prompt_message(max_token_limit=max_token_limit, massage_limit=massage_limit)
        #2、调用langchain 集成的get_buffer_string() 函数将消息体转换成文本
        return get_buffer_string(messages, human_prefix, ai_prefix)
